from args import parse_args
from data.datamodule import FundusDataModule
from models.main import ViTAD
import lightning as L
if __name__ == "__main__":
    args, cfg = parse_args()
    # Load Data Module
    data_module = FundusDataModule(
        data_dir=cfg['data']['data_dir'],
        batch_size=cfg['trainer']['batch_size'],
        input_size=cfg['data'].get('input_size', 224),
        num_workers=cfg['trainer'].get('num_workers', 4)
    )

    # Load Model
    model = ViTAD(
        model_t=cfg['model']['model_t'],
        model_f=cfg['model']['model_f'],
        model_s=cfg['model']['model_s'],
        opti_cfg=cfg['optimizer'],
        scheduler_cfg=cfg['scheduler'],
        freeze_t=cfg['model'].get('freeze_t', True),
        batch_size=cfg['trainer']['batch_size'], 
        loss_terms=cfg['model'].get('loss_terms', None),
        test_cfg=cfg.get('test_cfg', None)
        )

    # Trainer
    trainer = L.Trainer(
        max_epochs=cfg['trainer']['num_epochs'],
        accelerator=cfg['trainer']['accelerator'],
        precision=cfg['trainer']['precision'],
        logger=True,
        devices=cfg['trainer']['devices'],
        default_root_dir=cfg['trainer']['save_dir'],
    )

    # Train
    trainer.fit(model, datamodule=data_module)

    # Test
    trainer.test(model, datamodule=data_module)